#include "GMRESSolver.hpp"

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
GMRESSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::GMRESSolver(): IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>()
{
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void GMRESSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetRestartlen(int rstl)
{
    restart_len = rstl;
}

// 抄的hypre的GMRES，比较麻烦，精简了一点
template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void GMRESSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Mult_ipl(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const 
{
    assert(this->oper != nullptr);
    assert(this->restart_len > 0);

    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    if (my_pid == 0) printf("GMRES with restart = %d\n", restart_len);
    double * record = this->part_times;
    memset(record, 0, sizeof(double) * NUM_KRYLOV_RECORD);

    ksp_t * rs = new ksp_t [restart_len + 1];
    ksp_t *  c = new ksp_t [restart_len];
    ksp_t *  s = new ksp_t [restart_len];
    ksp_t * rs2= new ksp_t [restart_len + 1];
    ksp_t ** hh = new ksp_t * [restart_len + 1];
    for (int i = 0; i <= restart_len; i++)
        hh[i] = new ksp_t [restart_len];

    double t_allred = 0.0;

    // 初始化辅助向量
    par_Vector<idx_t, ksp_t, dof> r(x), w(x);
    r.set_val(0.0);
    w.set_val(0.0);
    std::allocator<par_Vector<idx_t, ksp_t, dof> > alloc;
    par_Vector<idx_t, ksp_t, dof> * p = alloc.allocate(restart_len + 1);
    for (int i = 0; i <= restart_len; i++) {
        alloc.construct(p + i, x);
        p[i].set_val(0.0);
    }
    par_Vector<idx_t, pc_calc_t, dof> * pc_buf_b = nullptr, * pc_buf_x = nullptr;
    if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
        pc_buf_b = new par_Vector<idx_t, pc_calc_t, dof>(b.comm, b.glb_nrows, b.beg_row, b.end_row);
        pc_buf_x = new par_Vector<idx_t, pc_calc_t, dof>(*pc_buf_b);
        pc_buf_b->set_val(0.0);
        pc_buf_x->set_val(0.0);
    }
    ksp_t b_norm = this->Norm(b);
    ksp_t r_norm, real_r_norm_old = b_norm, real_r_norm_new;

    // compute initial residual
    if (this->zero_guess) {
        vec_copy(b, p[0]);
        r_norm = b_norm;
    } else {
        record[OPER] -= wall_time();
        this->oper->Mult(x, p[0], false);
        record[OPER] += wall_time();
        vec_add(b, (ksp_t) -1.0, p[0], p[0]);
        r_norm = this->Norm(p[0]);
    }

    if (my_pid == 0) {
        printf("L2 norm of b: %20.16e\n", (double)b_norm);
        printf("Initial L2 norm of residual: %20.16e\n", (double)r_norm);
    }

    int & iter = this->final_iter = 0;
    ksp_t den_norm;
    if (b_norm > 0.0)   den_norm = b_norm;// convergence criterion |r_i|/|b| <= accuracy if |b| > 0
    else                den_norm = r_norm;// convergence criterion |r_i|/|r0| <= accuracy if |b| = 0

    /*  convergence criteria: |r_i| <= max( a_tol, r_tol * den_norm)
        den_norm = |r_0| or |b|
        note: default for a_tol is 0.0, so relative residual criteria is used unless
        user specifies a_tol, or sets r_tol = 0.0, which means absolute
        tol only is checked  */
    ksp_t epsilon;
    epsilon = std::max(this->abs_tol, this->rel_tol * den_norm);

    // so now our stop criteria is |r_i| <= epsilon
    if (my_pid == 0) {
        if (b_norm > 0.0) {
            printf("=========================================================\n\n");
            printf("Iters             resid.norm               rel.res.norm  \n");
            printf("------       -------------------      --------------------\n");
        } else {
            printf("=============================================\n\n");
            printf("Iters     resid.norm     \n");
            printf("------    ------------    \n");
        }
    }

    /* once the rel. change check has passed, we do not want to check it again */
    // bool rel_change_passed = false;
    while (iter < this->max_iter) {
        /* initialize first term of hessenberg system */
        rs[0] = r_norm;
        if (r_norm == 0.0) {
            this->converged = 1;
            this->final_norm = r_norm;
            goto finish;
        }

        /* see if we are already converged and 
           should print the final norm and exit */
        if (r_norm <= epsilon) {
            record[OPER] -= wall_time();
            this->oper->Mult(x, r, false);
            record[OPER] += wall_time();
            vec_add(b, (ksp_t) -1.0, r, r);
            r_norm = this->Norm(r);
            if (r_norm <= epsilon) {
                if (my_pid == 0)
                    printf("\nFinal L2 norm of residual: %20.16e\n\n", (double)r_norm);
                break;
            } else {
                if (my_pid == 0) 
                    printf("false convergence\n");
            }
        }

        vec_scale( (ksp_t)(1.0 / r_norm), p[0]);
        int i = 0;
        while (i < restart_len && iter < this->max_iter) {
            i++;
            iter++;

            if (this->prec) {
                bool zero_guess = true;
                if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                    VecTransPrec(p[i-1], *pc_buf_b);// copy in
                    if (zero_guess) pc_buf_x->set_val(0.0);
                    else
                        VecTransPrec(r, *pc_buf_x);// copy in
                    record[PREC] -= wall_time();
                    this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
                    record[PREC] += wall_time();
                    VecTransPrec(*pc_buf_x, r);
                } else {
                    if (zero_guess) r.set_val(0.0);
                    record[PREC] -= wall_time();
                    this->prec->Mult(p[i-1], r, zero_guess);
                    record[PREC] += wall_time();
                }
            }
            else                    vec_copy(p[i-1], r);

            record[OPER] -= wall_time();
            this->oper->Mult(r, p[i], false);
            record[OPER] += wall_time();

            /* modified Gram_Schmidt */
            for (int j = 0; j < i; j++) {
                t_allred -= wall_time();
                hh[j][i-1] = this->Dot(p[j], p[i]);
                t_allred += wall_time();
                vec_add(p[i], -hh[j][i-1], p[j], p[i]);
            }

            t_allred -= wall_time();
            ksp_t t = this->Norm(p[i]);
            t_allred += wall_time();
            hh[i][i-1] = t;	
            if (t != 0.0) vec_scale( (ksp_t)(1.0 / t), p[i]);

            /* done with modified Gram_schmidt and Arnoldi step.
                update factorization of hh */
            for (int j = 1; j < i; j++) {
                ksp_t t = hh[j-1][i-1];
                hh[j-1][i-1] = s[j-1]*hh[j][i-1] + c[j-1] * t;
                hh[j][i-1]   = -s[j-1]*t + c[j-1]*hh[j][i-1];
            }
            t = hh[i][i-1] * hh[i][i-1];
            t += hh[i-1][i-1] * hh[i-1][i-1];
            ksp_t gamma;
#ifdef __x86_64__
            if constexpr (std::is_same<ksp_t, __float128>::value)
                gamma = sqrtq(t);
            else
#endif
                gamma = sqrt(t);
            if (gamma == 0.0) gamma = 1.0e-16;
            c[i-1] = hh[i-1][i-1]/gamma;
            s[i-1] = hh[i][i-1]/gamma;
            rs[i] = -hh[i][i-1]*rs[i-1];
            rs[i]/=  gamma;
            rs[i-1] = c[i-1]*rs[i-1];
            /* determine residual norm */
            hh[i-1][i-1] = s[i-1]*hh[i][i-1] + c[i-1]*hh[i-1][i-1];
#ifdef __x86_64__
            if constexpr (std::is_same<ksp_t, __float128>::value)
                r_norm = fabsq(rs[i]);
            else
#endif
            r_norm = fabs(rs[i]);

            if (my_pid == 0) {
                if (b_norm > 0.0)
                    printf("%5d    %.16e   %.16e\n", iter, (double)r_norm, (double)(r_norm/b_norm));
                else
                    printf("%5d    %.16e\n", iter, (double)r_norm);
            }

            if (r_norm <= epsilon) break;
        }/*** end of restart cycle ***/

        /* now compute solution, first solve upper triangular system */
        rs[i-1] = rs[i-1]/hh[i-1][i-1];
        for (int k = i-2; k >= 0; k--)
        {
            ksp_t t = 0.0;
            for (int j = k+1; j < i; j++) {
                t -= hh[k][j]*rs[j];
            }
            t += rs[k];
            rs[k] = t/hh[k][k];
        }

        vec_mul_by_scalar(rs[i-1], p[i-1], w);
        for (int j = i-2; j >= 0; j--)
            vec_add(w, rs[j], p[j], w);

        if (this->prec) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                VecTransPrec(w, *pc_buf_b);// copy in
                if (zero_guess) pc_buf_x->set_val(0.0);
                else
                    VecTransPrec(r, *pc_buf_x);// copy in 
                record[PREC] -= wall_time();
                this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
                record[PREC] += wall_time();
                VecTransPrec(*pc_buf_x, r);
            } else {
                if (zero_guess) r.set_val(0.0);
                record[PREC] -= wall_time();
                this->prec->Mult(w, r, zero_guess);
                record[PREC] += wall_time();
            }
        }
        else                    vec_copy(w, r);

        /* update current solution x (in x) */
        vec_add(x, (ksp_t) 1.0, r, x);

        /* check for convergence by evaluating the actual residual */
        if (r_norm <= epsilon) {
            record[OPER] -= wall_time();
            this->oper->Mult(x, r, false);
            record[OPER] += wall_time();

            vec_add(b, (ksp_t) -1.0, r, r);
            real_r_norm_new = r_norm = this->Norm(r);
            this->final_norm = r_norm;

            if (r_norm <= epsilon) {
                if (my_pid == 0)
                    printf("\nFinal L2 norm of residual: %20.16e\n\n", (double)r_norm);
                goto finish;
            } else {/* conv. has not occurred, according to true residual */
                /* exit if the real residual norm has not decreased */
                if (real_r_norm_new >= real_r_norm_old) {
                    if (my_pid == 0)
                        printf("\nFinal L2 norm of residual: %20.16e breakdown real residual not decreased\n\n", (double)r_norm);
                    this->converged = 1;
                    break;
                }

                /* report discrepancy between real/GMRES residuals and restart */
                if (my_pid == 0)
                    printf("false convergence 2, L2 norm of residual: %20.16e\n", (double)r_norm);

                vec_copy(r, p[0]);
                i = 0;
                real_r_norm_old = real_r_norm_new;
            }
        }/* end of convergence check */

        /* compute residual vector and continue loop */
	    for (int j=i ; j > 0; j--) {
            rs[j-1] = -s[j-1]*rs[j];
            rs[j] = c[j-1]*rs[j];
    	}
        
        if (i) vec_add(p[i], (ksp_t) (rs[i] - 1.0), p[i], p[i]);

        for (int j=i-1 ; j > 0; j--)
            vec_add(p[i], rs[j], p[j], p[i]);
        
        if (i) {
            vec_add(p[0], (ksp_t) (rs[0] - 1.0), p[0], p[0]);
            vec_add(p[0], (ksp_t) 1.0, p[i], p[0]);
        }

        if (my_pid == 0) printf("Restarting...\n");
    }/* END of iteration while loop */

finish:
    delete [] c; delete [] s; delete [] rs; delete [] rs2;
    c = s = rs = rs2 = nullptr;
    for (int i = 0; i <= restart_len; i++) {
        delete [] hh[i]; hh[i] = nullptr;
    }
    delete [] hh; hh = nullptr;
    for (int i = 0; i <= restart_len; i++) alloc.destroy(p + i);
    alloc.deallocate(p, restart_len + 1);
    if (pc_buf_x) delete pc_buf_x;
    if (pc_buf_b) delete pc_buf_b;

    double min_allred = 0.0;
    MPI_Reduce(& t_allred, & min_allred, 1, MPI_DOUBLE, MPI_MIN, 0, comm);
    if (my_pid == 0) {
        printf("Proc %d OPER %.5f PREC %.5f AXPY %.5f DOT %.5f allred %.5f %.5f\n", my_pid,
            record[OPER], record[PREC], record[AXPY], record[DOT], t_allred, min_allred);
    }
}

template class GMRESSolver<int, double, double, double, 1>;
template class GMRESSolver<int, double, double, double, 2>;
template class GMRESSolver<int, double, double, double, 3>;
template class GMRESSolver<int, double, double, double, 4>;

template class GMRESSolver<int, float , double, double, 1>;
template class GMRESSolver<int, float , double, double, 2>;
template class GMRESSolver<int, float , double, double, 3>;
template class GMRESSolver<int, float , double, double, 4>;

template class GMRESSolver<int, float , float , double, 1>;
template class GMRESSolver<int, float , float , double, 2>;
template class GMRESSolver<int, float , float , double, 3>;
template class GMRESSolver<int, float , float , double, 4>;

template class GMRESSolver<int, float , float , float , 1>;
template class GMRESSolver<int, float , float , float , 2>;
template class GMRESSolver<int, float , float , float , 3>;
template class GMRESSolver<int, float , float , float , 4>;
#ifdef USE_FP16
template class GMRESSolver<int, __fp16, float , double, 1>;
template class GMRESSolver<int, __fp16, float , double, 2>;
template class GMRESSolver<int, __fp16, float , double, 3>;
template class GMRESSolver<int, __fp16, float , double, 4>;

template class GMRESSolver<int, __fp16, float , float , 1>;
template class GMRESSolver<int, __fp16, float , float , 2>;
template class GMRESSolver<int, __fp16, float , float , 3>;
template class GMRESSolver<int, __fp16, float , float , 4>;
#endif

